"""
使用streamlit作为web框架的回测结果查看器、参数优化结果查看器。
"""
import json
import pickle
import copy
import time

import pandas as pd
import streamlit as st
from streamlit_echarts import st_echarts
import os

st.set_page_config(layout="wide")


class DataLoader:
    def __init__(self, data_type):
        assert data_type in ('crypto', 'stock', 'futures'), "数据类型应该为 crypto, stock, futures三种！"
        self._type = data_type


def backtest(result_path):
    """
    对于回测结果的查看器
    """

    st.sidebar.title('回测策略查看器')
    options = ["index", "权益走势图", "K线图", "回测详细信息"]
    selected_tab = st.sidebar.radio("Choose a tab:", options)
    res = st.sidebar.file_uploader("选择回测数据！", accept_multiple_files=True)
    for file in res:  # TODO 需要变成配置项，不能把文件名写死。
        if file.name == "trades.csv":
            st.session_state['trade_record'] = pd.read_csv(file)
        elif file.name == "factor_feed.pkl":
            _ = pickle.load(file)
            _ = _.get_all()
            st.session_state["factor_feed"] = _
        elif file.name == "images_data.pkl":
            st.session_state['equity_data'] = pickle.load(file)
        elif file.name == 'plot_config.json':
            st.session_state['plot_config'] = json.load(file)
        elif file.name == 'qs_res.html':
            st.session_state['qs_res'] = file.read()

    # st.write("You selected:", selected_tab)
    p = os.path.dirname(os.path.abspath(__file__))
    if selected_tab == '权益走势图':

        with open(p + r'/backtest_echarts_option.json', 'r', encoding='utf-8') as f:
            option = json.load(f)
        # left, right = st.columns(2)
        # with left:
        #     uploaded_file = st.file_uploader("请选择要查看的pkl文件！")
        # with right:
            r = st.selectbox('', ('Log', 'Plain'))
            st.write('选择的策略是：')

        if 'equity_data' in st.session_state:
            res = st.session_state['equity_data']
            if 'stg_name' in res:
                st.write(res['stg_name'])

            # st.write(res)
            # with open(result_path + r'/images_data.pkl', 'rb') as f:
            # 策略应该存储一个字典内容为：{"equity": [...], "benchmark": [...], "draw_down": [...],
            # "equity_vs_market_value": [....], "datetime": [...], "equity_log": [....]}
            # res = pickle.load(f)
            # st.write(res)
            option['xAxis'][0]["data"] = res['datetime']
            option['xAxis'][1]["data"] = res['datetime']
            option['xAxis'][2]["data"] = res['datetime']
            option['series'][1]["data"] = res['draw_down']
            option['series'][2]["data"] = res['equity_vs_market_value']
            option['series'][3]["data"] = res['benchmark']
            if r == 'Log':
                option['series'][0]["data"] = res['equity_log']
            else:
                option['series'][0]["data"] = res['equity']
                option['yAxis'][0]['type'] = 'value'
            time.sleep(10)
            st_echarts(options=option, height="700px", width="100%")

    elif selected_tab == "K线图":
        if 'index' not in st.session_state:
            st.session_state['index'] = {"default": ""}
            for path in cfg['ploter']['data_path']:
                if path['path']:
                    for root, dir, files in os.walk(path['path']):
                        for file in files:
                            if files == '.DS_store':
                                continue
                            st.session_state['index'][file.split('.')[1] if path['type'] == 'futures' else file.split('.')[0]] = \
                                {
                                    "path": root + file,
                                    "type": path["type"]
                                 }

        idx = pd.Series(st.session_state['index'].keys())

        instrument = st.selectbox('选择查看的标的:', options=idx, index=0)
        if 'factor_feed' in st.session_state:
            factor_feed = st.session_state['factor_feed']
        else:
            factor_feed = None
        if 'trade_record' in st.session_state:
            trade_record = st.session_state['trade_record']
        else:
            trade_record = pd.DataFrame()

        with open(p + r'/backtest_kline_option.json', 'r') as f:
            option = json.load(f)
        if instrument != 'default':
            ins_info = st.session_state['index'][instrument]
            if os.path.exists(ins_info['path']):
                if ins_info['type'] == 'futures':  # 分不同标的来处理数据，主要是数据文件长得不一样
                    data = pd.read_csv(ins_info['path'])
                    option["series"][1]["data"] = data["volume"].tolist()
                    option["xAxis"][0]["data"] = data['date'].tolist()
                    option["xAxis"][1]["data"] = data['date'].tolist()
                    option["series"][0]["data"] = data.loc[:, ['open', 'close', 'low', 'high']].to_dict('split')['data']
                elif ins_info['type'] == 'stock':
                    data = pd.read_csv(ins_info['path'], encoding='gbk', skiprows=1)
                    option["series"][1]["data"] = data["成交量"].tolist()
                    option["xAxis"][0]["data"] = data['交易日期'].tolist()
                    option["xAxis"][1]["data"] = data['交易日期'].tolist()
                    option["series"][0]["data"] = data.loc[:, ['开盘价', '收盘价', '最低价', '最高价']].to_dict('split')['data']
                elif ins_info['type'] == 'crypto':
                    data = pd.DataFrame()
                else:
                    st.write(f'错误的标的类型：{ins_info["type"]}')
                    assert False
                if factor_feed and instrument in factor_feed and 'plot_config' in st.session_state:
                    # plot_config 是一个jsonconfig，外层是一个dict，包含要显示的factoer feed的字段，内层是这个字段的先显示是主图还是附图。
                    # {"ma": {"where":"main"},  # 表示主图
                    #     "cci": {"where": "auxiliary"}  # 表示副图
                    # }
                    series_auxiliary = {
                        "type": "line",
                        "name": "factor",
                        "xAxisIndex": 2,
                        "yAxisIndex": 2,
                        # "color": "gray",
                        "data": [],
                        "lineStyle": {
                            "width": 1
                            }
                        }
                    series_main = {
                        "type": "line",
                        "name": "factor",
                        "data": [],
                        "lineStyle": {
                            "width": 1
                        }
                    }
                    plot_config = st.session_state['plot_config']
                    # st.write(plot_config)
                    for k, v in plot_config.items():
                        # st.write(k)
                        if v['where'] == "main":
                            tp = copy.deepcopy(series_main)
                            tp["data"] = factor_feed[instrument][k].fillna(method='bfill').tolist()
                            tp['name'] = k
                            option["series"].append(tp)
                        elif v['where'] == "auxiliary":
                            tp = copy.deepcopy(series_auxiliary)
                            tp["data"] = factor_feed[instrument][k].fillna(method='bfill').tolist()
                            tp['name'] = k
                            option["series"].append(tp)
                            option["xAxis"][2]["data"] = data['date'].tolist()
                        else:
                            assert False, '未知的显示位置！'
                        option['legend']['data'].append(k)

            else:
                st.write('文件不存在，请重新输入！')
        if not trade_record.empty:
            point = {
                "coord": ["2023-05-25", 10],
                "symbol": "arrow",
                "symbolSize": 12,
                "symbolRotate": 180,
                "label": {"formatter": "123"},
                "itemStyle": {"color": "grey"}
            }
            for i, row in trade_record.iterrows():
                _ = copy.deepcopy(point)
                _['coord'] = [row['date'], row['price']]
                if row['action'] == "BUY":
                    _['label']['formatter'] = f"{row['volume']}"
                    _['symbolRotate'] = 0
                    _['itemStyle']['color'] = 'red'

                elif row['action'] == "SELL_SHORT":
                    _['label']['formatter'] = f"{row['volume']}"
                    _['itemStyle']['color'] = 'green'

                elif row['action'] == "SELL":
                    _['label']['formatter'] = f"{row['volume']}"
                elif row['action'] == "BUY_TO_COVER":
                    _['symbolRotate'] = 0
                    _['label']['formatter'] = f"{row['volume']}"
                option["series"][0]["markPoint"]["data"].append(_)
        st_echarts(options=option, height="700px", width="100%")
        a, b = st.columns(2)
        if not trade_record.empty:
            with a:
                st.dataframe(trade_record)
        if factor_feed and instrument in factor_feed:
            with b:
                tp = factor_feed[instrument].copy(deep=True)
                if 'time_delta_for_next_bar' in tp:
                    tp['time_delta_for_next_bar'] = str(factor_feed[instrument]['time_delta_for_next_bar'])
                st.dataframe(tp)
        # st.write(option)

    elif selected_tab == "回测详细信息":
        if 'qs_res' not in st.session_state:
            st.write(f"没有qs_res!")
        else:
            st.markdown(st.session_state['qs_res'], unsafe_allow_html=True)

    else:
        st.write("元亨利贞的策略查看器~")


if __name__ == '__main__':
    # backtest(r"/Users/yhlz/work_place/research/RB/rb_sma_new/RbSMA_short_118_long_66681b14ba",
    #          '/Users/yhlz/work_place/data/期货全部合约一分钟-2023-01-01/XSGE/RB')
    # backtest('D:/work/横盘突破策略/strategy/MinBreak_period_126_selected_num_200_volume_factor_2.5_',
    # 'D:\work\横盘突破策略\strategy\MinBreak_period_126_selected_num_200_volume_factor_2.5_\qs_res.html')
    with open(os.path.dirname(os.path.dirname(__file__)) + '/config/config.json') as f:
        cfg = json.load(f)
    backtest('')
